package org.hong.monkey.serializer

import java.io._
import java.nio.ByteBuffer

import org.hong.monkey.MonkeyConf
import org.hong.monkey.annotation.DeveloperApi
import org.hong.monkey.util.ByteBufferInputStream

import scala.reflect.ClassTag


private[monkey] class JavaSerializationStream(out: OutputStream, counterReset: Int, extraDebugInfo: Boolean)
  extends SerializationStream {
  private val objOut = new ObjectOutputStream(out)
  private var counter = 0

  def writeObject[T: ClassTag](t: T): SerializationStream = {
    try {
      objOut.writeObject(t)
    } catch {
      case e:NotSerializableException if extraDebugInfo =>
        throw SerializationDebugger.improveException(t, e)
    }
    counter += 1
    if (counterReset > 0 && counter < counterReset) {
      objOut.reset()
      counter = 0
    }
    this
  }

  def flush() { objOut.flush() }
  def close() { objOut.close() }
}

private[monkey] class JavaDeserializationStream(in: InputStream, loader: ClassLoader)
  extends DeserializationStream {

  private val objIn = new ObjectInputStream(in) {
    override def resolveClass(desc: ObjectStreamClass): Class[_] =
      try {
        Class.forName(desc.getName, false, loader)
      } catch {
        case e: ClassNotFoundException =>
          JavaDeserializationStream.primitiveMappings.get(desc.getName).getOrElse(throw e)
      }
  }

  def readObject[T: ClassTag](): T = objIn.readObject().asInstanceOf[T]
  def close() { objIn.close() }
}

private object JavaDeserializationStream {
  val primitiveMappings = Map[String, Class[_]](
  "boolean" -> classOf[Boolean],
  "byte" -> classOf[Byte],
  "char" -> classOf[Char],
  "short" -> classOf[Short],
  "int" -> classOf[Int],
  "long" -> classOf[Long],
  "float" -> classOf[Float],
  "double" -> classOf[Double],
  "void" -> classOf[Void]
  )
}

private[monkey] class JavaSerializerInstance(counterReset:Int, extraDebugInfo:Boolean, defaultClassLoader:ClassLoader)
  extends SerializerInstance {

  override def serialize[T: ClassTag](t: T): ByteBuffer = {
    val bos = new ByteArrayOutputStream()
    val out = serializeStream(bos)
    out.writeObject(t)
    out.close()
    ByteBuffer.wrap(bos.toByteArray)
  }

  override def deserialize[T: ClassTag](bytes: ByteBuffer): T = {
    val bis = new ByteBufferInputStream(bytes)
    val in = deserializeStream(bis)
    in.readObject()
  }

  override def deserialize[T: ClassTag](bytes: ByteBuffer, loader: ClassLoader): T = {
    val bis = new ByteBufferInputStream(bytes)
    val in = deserializeStream(bis, loader)
    in.readObject()
  }

  override def serializeStream(s: OutputStream): SerializationStream = {
    new JavaSerializationStream(s, counterReset, extraDebugInfo)
  }

  override def deserializeStream(s: InputStream): DeserializationStream = {
    new JavaDeserializationStream(s, defaultClassLoader)
  }

  def deserializeStream(s: InputStream, loader: ClassLoader): DeserializationStream = {
    new JavaDeserializationStream(s, loader)
  }
}

@DeveloperApi
class JavaSerializer(conf: MonkeyConf) extends Serializer with Externalizable {

  private var counterReset = conf.getInt("monkey.serializer.objectStreamReset", 100)
  private var extraDebugInfo = conf.getBoolean("monkey.serializer.extraDebugInfo", true)

  protected def this() = this(new MonkeyConf())

  override def newInstance(): SerializerInstance = {
    val classLoader = defaultClassLoader.getOrElse(Thread.currentThread().getContextClassLoader)
    new JavaSerializerInstance(counterReset, extraDebugInfo, classLoader)
  }

  override def writeExternal(out: ObjectOutput): Unit = {
    out.writeInt(counterReset)
    out.writeBoolean(extraDebugInfo)
  }

  override def readExternal(in: ObjectInput): Unit = {
    counterReset = in.readInt()
    extraDebugInfo = in.readBoolean()
  }

}
